import cv2
import numpy as np
import ffmpeg
from pathlib import Path
from typing import List, Union, Generator, Optional
import logging

logger = logging.getLogger(__name__)

class VideoProcessor:
    def __init__(self, batch_size: int = 100):
        """
        Initialize the video processor

        Args:
            batch_size: Number of frames to process at once
        """
        self.batch_size = batch_size

    def get_video_info(self, video_path: Union[str, Path]) -> dict:
        """
        Get video information

        Args:
            video_path: Path to the video file

        Returns:
            Dictionary containing video information
        """
        cap = cv2.VideoCapture(str(video_path))
        info = {
            'frame_count': int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),
            'fps': int(cap.get(cv2.CAP_PROP_FPS)),
            'width': int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            'height': int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        }
        cap.release()
        return info

    def frame_generator(self, video_path: Union[str, Path]) -> Generator[np.ndarray, None, None]:
        """
        Generate frames from video file one at a time

        Args:
            video_path: Path to the video file

        Yields:
            Video frames as numpy arrays
        """
        cap = cv2.VideoCapture(str(video_path))
        try:
            while cap.isOpened():
                ret, frame = cap.read()
                if not ret:
                    break
                yield frame
        finally:
            cap.release()

    def extract_frames(self, video_path: Union[str, Path], start_frame: Optional[int] = None, end_frame: Optional[int] = None) -> List[np.ndarray]:
        """
        Extract frames from a video file in batches

        Args:
            video_path: Path to the video file
            start_frame: Start frame number (optional)
            end_frame: End frame number (optional)

        Returns:
            List of frames as numpy arrays
        """
        frames = []
        frame_count = 0
        total_frames = 0

        # Get video info
        video_info = self.get_video_info(video_path)
        logger.info(f"Video info: {video_info}")

        for frame in self.frame_generator(video_path):
            if start_frame is not None and frame_count < start_frame:
                frame_count += 1
                continue

            if end_frame is not None and frame_count >= end_frame:
                break

            frames.append(frame)
            frame_count += 1
            total_frames += 1

            # Process in batches to save memory
            if len(frames) >= self.batch_size:
                logger.info(f"Processed {total_frames} frames out of {video_info['frame_count']}")
                yield frames
                frames = []

        # Yield remaining frames
        if frames:
            logger.info(f"Processed {total_frames} frames out of {video_info['frame_count']}")
            yield frames

    def extract_audio(self, video_path: Union[str, Path]) -> Path:
        """
        Extract audio from a video file

        Args:
            video_path: Path to the video file

        Returns:
            Path to the extracted audio file
        """
        video_path = Path(video_path)
        audio_path = video_path.parent / f"{video_path.stem}_audio.wav"

        stream = ffmpeg.input(str(video_path))
        stream = ffmpeg.output(stream, str(audio_path))
        ffmpeg.run(stream, overwrite_output=True)

        return audio_path

    def combine_video_audio(
        self,
        video_path: Union[str, Path],
        audio_path: Union[str, Path],
        output_path: Union[str, Path],
        fps: int = 30
    ) -> Path:
        """
        Combine video and audio into a final video

        Args:
            video_path: Path to the input video file
            audio_path: Path to the audio file
            output_path: Path for the output video
            fps: Frames per second

        Returns:
            Path to the output video
        """
        output_path = Path(output_path)

        # Combine video with audio
        video = ffmpeg.input(str(video_path))
        audio = ffmpeg.input(str(audio_path))
        stream = ffmpeg.output(
            video,
            audio,
            str(output_path),
            vcodec='libx264',
            acodec='aac'
        )
        ffmpeg.run(stream, overwrite_output=True)

        # Verify the output file exists
        if not output_path.exists():
            raise RuntimeError(f"Failed to create output video: {output_path}")

        return output_path